#include "detect_client.h"

 /*
	预处理的缩放比例
	在不丢失原图比例的同时，尽可能的伸缩；同时为了保证检测效果，只允许缩放，不允许放大。
*/
float get_max_scale(int input_width, int input_height, int net_width, int net_height)
{
    float scale = min((float)input_width / net_width, (float)input_height/net_height);
    if(scale > 1) return 1;
    else return scale;
}
/*
	Yolov5 输出结果解码4 从decode.cpp搬运过来
	1.相对尺寸转回绝对尺寸 
	2.输出结果解码1-3完成了针对于深度学习网络尺寸的解码此时图片大小
	为640x640(maybe)
	由于图像可能经过预处理缩放过 因此还需进一步还原至原尺寸
*/
void reduction_size(vector<detection> &dets_perframe,cv::Mat &img,float scale,detection* dets,int nboxes_left,float thresh)
{
	for(int i=0;i<nboxes_left;i++){
		// char labelstr[4096]={0};
		detection t;
		t.sort_class = -1;
		if(dets[i].objectness==0) continue;
		for(int j=0;j<nclasses;j++){
			if(dets[i].prob[j]>thresh){
				t.sort_class = j;
				t.prob[j] = dets[i].prob[j];
				break;
			}
		}
		//如果t.sort_class>0说明框中有物体
		if(t.sort_class>=0){
			cv::Rect_<float> b=dets[i].bbox;
			//计算坐标 先根据缩放后的图计算绝对坐标 然后除以scale缩放到原来的图
			//又因为原点重合 因此缩放后的结果就是原结果
			int x1 = b.x * NET_INPUTWIDTH / scale;
			int x2= x1 + b.width * NET_INPUTWIDTH / scale + 0.5;
			int y1= b.y * NET_INPUTHEIGHT / scale;
			int y2= y1 + b.height * NET_INPUTHEIGHT / scale + 0.5;

            if(x1  < 0) x1  = 0;
            if(x2> img.cols-1) x2 = img.cols-1;
            if(y1 < 0) y1 = 0;
            if(y2 > img.rows-1) y2 = img.rows-1;

			t.bbox = cv::Rect_<float>(x1, y1, x2-x1, y2-y1);
			dets_perframe.push_back(t);
            }
		}
	return;
}

void get_detections(int cpuid, int sock_tcp[]){
	cpu_set_t mask;

	CPU_ZERO(&mask);
	CPU_SET(cpuid, &mask);

	if (pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask) < 0)
		cerr << "set thread affinity failed" << endl;

	printf("Bind get_detections process to CPU %d\n", cpuid); 

	int FINISH = 0;	// 循环跳出条件 所有的SOCEKT都半关闭
	float scale = get_max_scale(IMG_WIDTH, IMG_HEIGHT, NET_INPUTWIDTH, NET_INPUTHEIGHT); // 预处理缩放比例
	while(1){
		int idx, nboxes_left;
		int sock_idx = (idxOutputImage / BATCHSIZE) % SOCKETNUM;
		// 1 读idx
		int cur = 0, tmp = 0, length = sizeof(int);
		cur = read(sock_tcp[sock_idx], &idx, length);
		if(cur == 0){
			if(++FINISH == SOCKETNUM){
				printf("video_dect over.\n");
				break;
			}
			else{
				continue;
			}
		}
		else if(cur == -1){
			printf("video_dect error.\n");
			break;
		}
		else if(cur > 0 && cur < length){
			while(cur < length){
				if((tmp = read(sock_tcp[sock_idx], &idx+cur, length-cur)) <= 0){
					printf("detect_client.cpp:read idx error.\n");
					return;
				}
				cur += tmp;
			}
		}
		// 2 读NMS后剩下的预测框个数
		cur = 0;
		while(cur < length){
			if((tmp = read(sock_tcp[sock_idx], &nboxes_left+cur, length-cur)) <= 0){
				printf("detect_client.cpp:read nboxes_left error.\n");
				return;
			}
			cur += tmp;
		}
		// 3 读出每个预测框
		cur = 0;
		detection* dets=(detection*) calloc(nboxes_total,sizeof(detection));
		int total = sizeof(detection)*nboxes_left;
		while(cur < total){
			if((tmp = read(sock_tcp[sock_idx], dets+cur, total-cur)) <= 0){
				printf("detect_client.cpp:read detection* dets error.\n");
				delete dets;
				return;
			}
			cur += tmp;
		}
		if(queueInput_client.front().first == idx){
			cv::Mat img_src = queueInput_client.front().second;
			// draw_image(img_src, 1.0, dets, nboxes_left, DRAW_CLASS_THRESH);
			// imshow("display", img_src);
			// waitKey(20);
			imageout_idx t;
			t.frame_id = idx;
			t.img = img_src;
			vector<detection> tmp;
			reduction_size(tmp, img_src, scale, dets, nboxes_left, DRAW_CLASS_THRESH);
			t.dets = tmp;

			mtxQueueShow.lock();
			queueShow.push(t);
			mtxQueueShow.unlock();
			printf("img(%d) res: %d\n", idx, nboxes_left);
			mtxQueueInput_client.lock();
			queueInput_client.pop();
			mtxQueueInput_client.unlock();
		}
		else{
			printf("Disorder\n");
			break;
		}
		idxOutputImage++;
		delete dets;
	}
	bDetecting = false;
}

string labels[2]={"person", "vehicle"};
cv::Scalar colorArray[2]={
	cv::Scalar(139,0,0,255),
	cv::Scalar(139,0,139,255),
};
/*---------------------------------------------------------
	绘制预测框
----------------------------------------------------------*/
int draw_image(cv::Mat img,float scale,detection* dets,int total,float thresh){
	//::cvtColor(img, img, cv::COLOR_RGB2BGR);
	for(int i=0;i<total;i++){
		char labelstr[4096]={0};
		int class_=-1;
		int topclass=-1;
		float topclass_score=0;
		if(dets[i].objectness==0) continue;
		for(int j=0;j<nclasses;j++){
			if(dets[i].prob[j]>thresh){
				if(topclass_score<dets[i].prob[j]){
					topclass_score=dets[i].prob[j];
					topclass=j;
				}
				if(class_<0){
					strcat(labelstr,labels[j].data());
					class_=j;
				}
				else{
					strcat(labelstr,",");
					strcat(labelstr,labels[j].data());
				}
			}
		}
		//如果class>0说明框中有物体,需画框
		if(class_>=0){
			cv::Rect_<float> b=dets[i].bbox;
			//计算坐标 先根据缩放后的图计算绝对坐标 然后除以scale缩放到原来的图
			//又因为原点重合 因此缩放后的结果就是原结果
			int x1 = b.x * NET_INPUTWIDTH / scale;
			int x2= x1 + b.width * NET_INPUTWIDTH / scale + 0.5;
			int y1= b.y * NET_INPUTWIDTH / scale;
			int y2= y1 + b.height * NET_INPUTHEIGHT / scale + 0.5;

            if(x1  < 0) x1  = 0;
            if(x2> img.cols-1) x2 = img.cols-1;
            if(y1 < 0) y1 = 0;
            if(y2 > img.rows-1) y2 = img.rows-1;
			//std::cout << labels[topclass] << "\t@ (" << x1 << ", " << y1 << ") (" << x2 << ", " << y2 << ")" << "\n";

            rectangle(img, cv::Point(x1, y1), cv::Point(x2, y2), colorArray[class_%10], 3);
            putText(img, labelstr, cv::Point(x1, y1 - 12), 1, 2, cv::Scalar(0, 255, 0, 255));
            }
		}
	return 0;
}